package com.typhoon.spring_shiro.utils.limit;

import java.util.ArrayList;
import java.util.List;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.util.Assert;

import com.typhoon.spring_shiro.utils.context.SpringContextUtil;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;

/**
 * 分布式速率限制 例如:限制n秒钟请求x次
 * 
 * @author Typhoon
 * @date 2017-06-08 15:08
 * @since V2.0
 */
public class RateLimiter {

    // private static final org.slf4j.LOGGER logger = LoggerFactory.getLogger(RateLimiter.class);

    private static JedisPool jedisPool = SpringContextUtil.getBean("jedisPool",JedisPool.class);

    /**
     * 针对资源key,每seconds秒最多访问maxCount次,超过maxCount次返回false
     *
     * @param key
     * @param seconds
     * @param limitCount
     * @return
     */
    public static  boolean tryAccess(String key, int seconds, int limitCount) {
        LimitRule limitRule = new LimitRule();
        limitRule.setLimitCount(limitCount);
        limitRule.setSeconds(seconds);
        return tryAccess(key, limitRule);
    }

    /**
     * 针对资源key,每limitRule.seconds秒最多访问limitRule.limitCount,超过limitCount次返回false 超过lockCount
     * 锁定lockTime
     * 
     * @param key
     * @param limitRule
     * @return
     */
    public static boolean tryAccess(String key, LimitRule limitRule) {

        Jedis jedis = null;
        long count = -1;
        try {
            jedis = jedisPool.getResource();
            List<String> keys = new ArrayList<String>();
            String newKey = "limit:" + key + ":" + limitRule.getSeconds();
            keys.add(newKey);
            List<String> args = new ArrayList<String>();
            args.add(Math.max(limitRule.getLimitCount(), limitRule.getLockCount()) + "");
            args.add(limitRule.getSeconds() + "");
            args.add(limitRule.getLockCount() + "");
            args.add(limitRule.getLockTime() + "");

            count = Long.parseLong(jedis.eval(buildLuaScript(limitRule), keys, args) + "");
            return count <= limitRule.getLimitCount();
        } finally {
            if (jedis != null)
                jedis.close();
        }
    }

    /**
     * 针对多条规则进行验证
     * 
     * @param key
     * @param rules
     * @return
     */
    public static boolean tryAccess(String key, LimitRuleSet rules) {
        Assert.isTrue(rules != null && rules.size() > 0, "param 'rules' can not be empty");

        List<LimitRule> list = rules.list();
        for (int i = 0; i < list.size(); i++) {
            // 遍历全部规则，确保每个规则都有计数
            LimitRule rule = list.get(i);
            if (!tryAccess(key, rule)) {
                return false; // 频率超限
            }
        }

        return true;

    }

    /**
     * 创建lua脚本交给redis执行
     * <p>新版本redis支持lua脚本</p>
     * <p>lua脚本被执行的时候能够保证原子性</p>
     * @param limitRule
     * @return
     */
    private static String buildLuaScript(LimitRule limitRule) {
        StringBuilder lua = new StringBuilder();
        lua.append("\nlocal c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn c;");
        lua.append("\nend");
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        if (limitRule.enableLimitLock()) {
            lua.append("\nif tonumber(c) > tonumber(ARGV[3]) then");
            lua.append("\nredis.call('expire',KEYS[1],ARGV[4])");
            lua.append("\nend");
        }
        lua.append("\nreturn c;");
        return lua.toString();
    }

}
